package middleware

import (
	"github.com/gin-gonic/gin"
	"github.com/sirupsen/logrus"
	"net/http"
	"smart-sensor-cloud/pkg/auth"
	"smart-sensor-cloud/pkg/context"
	"smart-sensor-cloud/pkg/log"
	"smart-sensor-cloud/pkg/util"
	"strings"
)

const (
	TokenHeaderKey  = "Authorization"
	JWTGeneratorKey = "JWT_GENERATOR"
	UserClaimsKey   = "CUR_USER_CLAIMS"
	UserIDKey       = "__USER_ID__"
)

func JWTGeneratorMiddleware(jwt *auth.JWTGenerator) gin.HandlerFunc {
	return func(ctx *gin.Context) {
		ctx.Set(JWTGeneratorKey, jwt)
		ctx.Next()
	}
}

func wrapContextWithUserID(c *gin.Context, userID int) {
	c.Set(UserIDKey, userID)
	ctx := context.WithUserID(c.Request.Context(), userID)
	ctx = log.NewUserIDContext(ctx, userID)
	c.Request = c.Request.WithContext(ctx)
}

func JWTAuthMiddleware(skippers ...Skipper) gin.HandlerFunc {
	return func(ctx *gin.Context) {
		if SkipHandler(ctx, skippers...) {
			ctx.Next()
			return
		}

		token := GetToken(ctx)
		if token == "" {
			ctx.JSON(http.StatusOK, &util.Error{
				Code:    http.StatusUnauthorized,
				Message: "请求未携带Token，无法进行访问",
			})
			ctx.Abort()
			return
		}

		logrus.Infof("本次请求携带Token: %s", token)

		jwt := ctx.MustGet(JWTGeneratorKey).(*auth.JWTGenerator)

		claims, err := jwt.ParseTokenToClaims(token)
		if err != nil {
			ctx.JSON(http.StatusOK, &util.Error{
				Code:    http.StatusUnauthorized,
				Message: err.Error(),
			})
			ctx.Abort()
			return
		}

		wrapContextWithUserID(ctx, claims.UserID)

		ctx.Set(UserClaimsKey, claims)
		ctx.Next()
	}
}

func GetUserID(ctx *gin.Context) int {
	if val, ok := ctx.Get(UserIDKey); ok {
		if num, ok := val.(int); ok {
			return num
		}
	}
	return -1
}

func GetToken(ctx *gin.Context) string {
	var token string
	authHeader := ctx.GetHeader("Authorization")
	prefix := "Bearer "
	if authHeader != "" && strings.HasPrefix(authHeader, prefix) {
		token = authHeader[len(prefix):]
	}
	return token
}
